/* 树结构的访问和操作 */
import { remove } from 'lodash-es'

// 遍历树
type TraverseInfo = {
  node: object // 节点对象
  index: number // 在父节点的子元素数组中的索引
  level: number // 层级
  parent: object // 父节点
  list: object[] // 父节点的子元素数组
  count: number // 父节点的子元素数组的长度
}

// 遍历（深度优先 DFS）
function traverse(rootNodeOrList: any, visit: (item: any) => void, childrenKey = 'children') {
  const rootIsList = Array.isArray(rootNodeOrList)
  const rootNode = rootIsList ? wrapListToNode(rootNodeOrList, childrenKey) : rootNodeOrList
  _traverse(rootNode)

  function _traverse(node: any, level = 0) {
    const list = node[childrenKey]
    if (list == null) return
    list.forEach((item: any, index: number) => {
      const info: TraverseInfo = {
        node: item,
        index,
        level,
        parent: rootIsList && level === 0 ? null : node,
        list,
        count: list.length,
      }
      visit(info)
      if (item[childrenKey]) {
        _traverse(item, level + 1)
      }
    })
  }
}

// 遍历过滤
function traverseFilter(rootNodeOrList: any, tester: (item: any) => boolean, childrenKey = 'children') {
  const rootIsList = Array.isArray(rootNodeOrList)
  const rootNode = rootIsList ? wrapListToNode(rootNodeOrList, childrenKey) : rootNodeOrList
  _traverseFilter(rootNode, tester)

  function _traverseFilter(node: any, tester: (item: any) => boolean) {
    const list = node[childrenKey]
    if (list == null) return

    remove(list, (item: any) => !tester(item))

    list.forEach((item: any) => _traverseFilter(item, tester))
  }
}

// 列表封装为树节点
function wrapListToNode(list: any[], childrenKey: string) {
  return { [childrenKey]: list }
}

// 导出
export const tree = {
  traverse,
  traverseFilter,
}
